from typing import Callable, Tuple, Dict, Iterable, Any

import networkx as nx

from .baseparser import Parser


def find_and_merge_cycles(graph: nx.DiGraph,
                          merge_func: Callable[[nx.DiGraph, Iterable[str]], Tuple[str, Dict[str, Any]]]) -> nx.DiGraph:
    # 1. 找到所有强连通分量(循环)
    strongly_connected_components = list(nx.strongly_connected_components(graph))

    # 2. 创建一个新图用于存储结果
    merged_graph = graph.copy()

    # 3. 遍历每个强连通分量
    for scc in strongly_connected_components:
        # 如果强连通分量大小大于1，说明存在循环
        if len(scc) > 1:
            # 将该组件中的所有节点名称合并成一个新名称
            # new_node = '+'.join(sorted(str(n) for n in scc))

            new_node, new_node_data = merge_func(graph, scc)

            # 获取所有指向循环中任何节点的边
            in_edges = []
            for n in scc:
                in_edges.extend(merged_graph.in_edges(n, data=True))

            # 获取所有从循环中任何节点出发的边
            out_edges = []
            for n in scc:
                out_edges.extend(merged_graph.out_edges(n, data=True))

            # 删除循环中的所有节点
            merged_graph.remove_nodes_from(scc)

            # 添加合并后的新节点
            merged_graph.add_node(new_node, **new_node_data)

            # 重新添加入边和出边
            for u, v, data in in_edges:
                if u not in scc:  # 只添加来自循环外的边
                    merged_graph.add_edge(u, new_node, **data)

            for u, v, data in out_edges:
                if v not in scc:  # 只添加指向循环外的边
                    merged_graph.add_edge(new_node, v, **data)

    return merged_graph


def get_adjacency(graph: nx.DiGraph, node):
    dependencies = list(graph.successors(node))
    return {node: graph.nodes[node] for node in dependencies}


def merge_example(graph, scc):
    new_node = '+'.join(sorted(str(n) for n in scc))
    return new_node, {}

# 使用示例
def example():
    # 创建一个有向图
    G = nx.DiGraph()

    # 添加边以形成循环
    G.add_edges_from([
        (1, 2), (2, 3), (3, 1),  # 循环 1-2-3
        (3, 5), (1, 5), (1, 4), (2, 4),
        (3, 4), (4, 5), (5, 6),  # 非循环路径
        (6, 7), (7, 6)  # 另一个循环 6-7
    ])

    # 找到并合并循环
    merged_G = find_and_merge_cycles(G, merge_example)
    p = Parser()
    p.graph = G
    p.visualize('ori')
    p.graph = merged_G
    p.visualize('merged')

if __name__ == '__main__':
    example()